import sys
import numpy as np
from PyQt5.QtWidgets import QApplication, QMainWindow, QTreeView, QVBoxLayout, QWidget
from PyQt5.QtGui import QStandardItemModel, QStandardItem

# 示例数据
RESULT = {
    "NODE_RESULT": {
        "U1": np.array([1, 4, 7]),
        "U2": np.array([2, 5, 8]),
        "U3": np.array([3, 6, 9]),
        "Umag": np.sqrt(np.array([1, 4, 7]) ** 2 + np.array([2, 5, 8]) ** 2 + np.array([3, 6, 9]) ** 2),
        "RF1": np.array([10, 13, 16]),
        "RF2": np.array([11, 14, 17]),
        "RF3": np.array([12, 15, 18]),
        "RFmag": np.sqrt(np.array([10, 13, 16]) ** 2 + np.array([11, 14, 17]) ** 2 + np.array([12, 15, 18]) ** 2),
    },
    "ELEMENT_RESULT": {},
    "INTEGRATION_POINT_RESULT": {}
}

class ExampleApp(QMainWindow):
    def __init__(self):
        super().__init__()
        self.initUI()

    def initUI(self):
        self.setWindowTitle("QTreeView 示例")
        self.setGeometry(100, 100, 600, 400)

        # 创建 QTreeView 控件
        self.treeView = QTreeView(self)

        # 创建模型
        self.model = QStandardItemModel()
        self.model.setHorizontalHeaderLabels(["名称", "值"])

        # 构建树形结构
        self.buildTree(RESULT, self.model.invisibleRootItem())

        # 设置模型
        self.treeView.setModel(self.model)

        # 设置布局
        layout = QVBoxLayout()
        layout.addWidget(self.treeView)

        # 创建一个中心窗口
        centralWidget = QWidget(self)
        centralWidget.setLayout(layout)
        self.setCentralWidget(centralWidget)

    def buildTree(self, data, parentItem):
        for key, value in data.items():
            if isinstance(value, dict):
                # 如果值是字典,创建一个父节点
                parent = QStandardItem(key)
                parent.setEditable(False)
                parentItem.appendRow([parent, QStandardItem("")])
                self.buildTree(value, parent)
            else:
                # 如果值是数组或其他类型,创建子节点
                item = QStandardItem(key)
                item.setEditable(False)
                valueItem = QStandardItem(str(value))
                valueItem.setEditable(False)
                parentItem.appendRow([item, valueItem])

if __name__ == "__main__":
    app = QApplication(sys.argv)
    ex = ExampleApp()
    ex.show()
    sys.exit(app.exec_())